# -*- codeing = utf-8 -*-
# @Time : 2022/11/7 16:21
# @Author : xiaow
# @File : videosocket.py
# @Software : PyCharm
# 获取图片数据，每帧传递给前台从而做成视频


import asyncio
import time

import websockets
import base64
import cv2
import numpy as np
from concurrent.futures.thread import ThreadPoolExecutor
from multiprocessing import Process, Lock
import mysqldao
from websockets.exceptions import ConnectionClosedError

lock = Lock()
outlock = Lock()
infopics = []
videopics = []
startflag = False
videoname = 'xw.avi'
# 线程池
pool = ThreadPoolExecutor(max_workers=50)
# 读取视频流
capture = cv2.VideoCapture('rtsp://127.0.0.1/test')
if not capture.isOpened():
    print('quit')
    quit()
ret, frame = capture.read()
task_list = []
startflag=False

# 存储图片和信息的类
class MImg:
    #  t 时间  data 图像数据
    def __init__(self, t, data, code=-1, data2="", info=""):
        self.t = t
        self.data = data
        # 状态码 -1 没有检测图片  0正常  1裂纹
        self.code = code
        # 检测后的图片
        self.data2 = data2
        # 监测信息
        self.info = info


# 输出视频
def outVideo(out, frame):
    global startflag
    if startflag:
        out.write(frame)



# 处理图像 这里的img就是三通道图像数据
def dealwithpic(img, lock):
    global infopics
    # 操作图像start
    # 操作end
    lock.acquire()
    infopics.append(MImg(time.time(), img, 0, img, '检测正常'))
    infopics = sorted(infopics, key=lambda e: e.t)
    lock.release()


# 获取图像
def getpics():
    global startflag, videoname
    fps = 30
    size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(capture.get(cv2.CAP_PROP_FRAME_EHIGHT)))
    fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', '2')
    createflag = False
    global outlock
    global ret, frame, infopics, outpics, lock
    ret, frame = capture.read()
    all = 0
    while ret:
        # 建立连接之后开始读取视频
        # print(startflag)
        if not createflag:
            out = cv2.VideoWriter(videoname, fourcc, fps, size)
            createflag = True
        task_list.append(pool.submit(outVideo, out, frame))
        # all 是获取的帧数 用于判断是否读取该帧图片
        all += 1
        frame = cv2.flip(frame, -1)
        frame = cv2.flip(frame, 0)
        result, imgencode = cv2.imencode('.jpg', frame)
        data = np.array(imgencode)
        data = cv2.imdecode(data, cv2.IMREAD_COLOR)
        # 判断infopics中有无处理好的图像
        lock.acquire()
        if infopics.__len__() > 0:
            infopic = infopics.pop()
            videopics.append(MImg(time.time(), data, infopic.code, infopic.data, infopic.info))
        else:
            videopics.append(MImg(time.time(), data))
        lock.release()
        # 进行检测  30帧检测一次
        if all % 30 == 0:
            task_list.append(pool.submit(dealwithpic, data, lock))
            all = 0
        ret, frame = capture.read()


# 清理线程池中执行完的线程
def clearpool():
    while (True):
        if len(task_list) == 0:
            break
        for i in task_list:
            if i.done():
                task_list.remove(i)


# 向服务器端实时发送视频截图
async def send_msg(websocket):
    global lock, videopics
    while True:
        if videopics.__len__() > 0:
            # 加锁拿数据
            lock.acquire()
            videopics = sorted(videopics, key=lambda e: e.t)
            a = videopics.pop()
            lock.release()
            img_encode = cv2.imencode('.jpg', a.data)[1]
            base64_data = base64.b64encode(img_encode)
            base64_data_str = str(base64_data, encoding="utf-8")
            # code!=-1就代表有检测完成的图片
            if a.code != -1:
                img_encode2 = cv2.imencode('.jpg', a.data2)[1]
                base64_data2 = base64.b64encode(img_encode2)
                base64_data_str2 = str(base64_data2, encoding="utf-8")
            else:
                base64_data_str2 = ""
            result = '{"data":"%s","code":"%s","data2":"%s","info":"%s"}' % ("data:image/jpg;base64," + base64_data_str,
                                                                             a.code.__str__(),
                                                                             "data:image/jpg;base64," + base64_data_str2,
                                                                             a.info
                                                                             )
            try:
                await websocket.send(result)
            except ConnectionClosedError as e:
                break


async def serverRecv(websocket):
    global startflag, videoname, infopics, videopics
    # 设置存储名
    videoname = "D:\\new.avi"
    startflag = True
    while True:
        try:
            recv_text = await websocket.recv()
        except ConnectionClosedError as e:
            pass

        await send_msg(websocket)
        videopics.clear()
        infopics.clear()
        if startflag:
            startflag=False


if __name__ == '__main__':
    task_list.append(pool.submit(clearpool))
    task_list.append(pool.submit(getpics))
    server = websockets.serve(serverRecv, "127.0.0.1", "8888")
    asyncio.get_event_loop().run_until_complete(server)
    asyncio.get_event_loop().run_forever()
